Add with_rolled_lon to downscaling models#1237
Draft
frodre wants to merge 1 commit into
Draft
Conversation
Let models re-express their grid in a seam-crossing coarse domain's longitude convention while sharing network weights: - DiffusionModel.with_rolled_lon rebuilds the model through its constructor with full_fine_coords and static_inputs rolled to match the coarse grid. The roll is anchored on the western coarse-cell edge so the fine grid stays aligned to whole coarse cells. Returns self when no roll is needed. - DenoisingMoEPredictor.with_rolled_lon rolls every expert (preserving the shared-grid invariant) and rebuilds so the sigma dispatcher is reconstructed from the rolled experts. Adds tests for no-roll passthrough, coord shifting with shared weights, idempotency, coarse-cell alignment, and rolling all MoE experts.
This was referenced Jun 6, 2026
frodre
added a commit
that referenced
this pull request
Jun 8, 2026
…1234) First in a 5-PR stack adding support for longitude domains that cross the 0/360 prime meridian in downscaling. This standalone hardening PR moves expert grid-compatibility validation into the predictor constructor so every construction path is protected, not just the config-build path: only the primary expert's coordinates are used for input prep and output coords, so an expert built on a mismatched grid would otherwise silently downscale onto the wrong grid. Changes: - `fme.downscaling.predictors.serial_denoising`: move `_validate_experts_compatible` from `DenoisingMoEConfig.build` into `DenoisingMoEPredictor.__init__`, so it holds for `build`, `from_state`, and future callers (e.g. `with_rolled_lon`). - `fme.downscaling.test_models`: add `test_denoising_moe_predictor_rejects_mismatched_expert_grids`, constructing the predictor directly with mismatched-grid experts and asserting it raises. - [x] Tests added - [ ] If dependencies changed, "deps only" image rebuilt and "latest_deps_only_image.txt" file updated Base: `main` ### Stack | PR | Head → Base | Title | |----|-------------|-------| | [#1234](#1234) | `refactor/moe-validate-experts-init` → `main` | Validate expert grid compatibility in `DenoisingMoEPredictor.__init__` | | [#1235](#1235) | `feature/lon-roll-primitives` → PR1 | Add longitude roll primitives | | [#1236](#1236) | `feature/lon-roll-data-layer` → PR2 | Roll seam-crossing longitudes in the data layer | | [#1237](#1237) | `feature/lon-roll-model` → PR3 | Add with_rolled_lon to models | | [#1238](#1238) | `feature/lon-roll-integration` → PR4 | Roll the model in inference/predict/evaluator |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
PR 4 of 5 in the prime-meridian longitude stack. Lets a model re-express its grid in a seam-crossing coarse domain's longitude convention while sharing the trained network weights, so a single checkpoint can generate over a domain expressed west of 0 or east of 360. Depends on the roll primitives (PR 2) and the constructor-time expert validation (PR 1).
Changes:
fme.downscaling.models.DiffusionModel.with_rolled_lon: rebuild the model through its constructor withfull_fine_coordsandstatic_inputsrolled to match the coarse grid, anchored on the western coarse-cell edge so the fine grid stays aligned to whole coarse cells; returnsselfwhen no roll is needed.fme.downscaling.predictors.serial_denoising.DenoisingMoEPredictor.with_rolled_lon: roll every expert (preserving the shared-grid invariant) and rebuild so the sigma dispatcher is reconstructed from the rolled experts.fme.downscaling.test_models: tests for no-roll passthrough, coord shifting with shared weights, idempotency, coarse-cell alignment, and rolling all MoE experts.Base:
feature/lon-roll-data-layer(PR 3)Stack
refactor/moe-validate-experts-init→mainDenoisingMoEPredictor.__init__feature/lon-roll-primitives→ PR1feature/lon-roll-data-layer→ PR2feature/lon-roll-model→ PR3feature/lon-roll-integration→ PR4